use clap::Parser;
use scrap::{Capturer, Display};
use std::fs::File;
use std::io::{BufWriter, ErrorKind};
use std::process::Command;
use std::thread;
use std::time::{Duration, Instant};
use chrono::Local;
use gif::{Encoder, Frame, Repeat};
use image::RgbaImage;

/// 简单的电脑录屏工具
#[derive(Parser, Debug)]
#[command(author, version, about, long_about = None)]
struct Args {
    /// 录制时长（秒）
    #[arg(short, long, default_value_t = 5)]
    duration: u64,

    /// 输出文件路径（不指定则使用时间戳命名）
    #[arg(short, long)]
    output: Option<String>,

    /// 帧率（每秒帧数）
    #[arg(short, long, default_value_t = 10)]
    fps: u32,

    /// 显示器索引（0表示主显示器）
    #[arg(short = 'm', long, default_value_t = 0)]
    monitor: usize,

    /// 输出格式 (gif, png-sequence 或 mp4)
    #[arg(long, default_value = "gif")]
    format: String,

    /// 保存 MP4 后是否保留临时 PNG 文件
    #[arg(long, default_value_t = false)]
    keep_frames: bool,
}

fn main() {
    let args = Args::parse();

    println!("🎬 屏幕录制工具启动");
    println!("⏱️  录制时长: {} 秒", args.duration);
    println!("🎞️  帧率: {} FPS", args.fps);
    println!("📝 输出格式: {}", args.format);
    
    // 获取显示器
    let display = if args.monitor == 0 {
        Display::primary().expect("无法获取主显示器")
    } else {
        let displays = Display::all().expect("无法获取显示器列表");
        if args.monitor >= displays.len() {
            eprintln!("❌ 显示器索引超出范围！可用显示器数量: {}", displays.len());
            return;
        }
        displays.into_iter().nth(args.monitor).expect("无法获取指定显示器")
    };
    
    println!("🖥️  使用显示器: {} ({}x{})", args.monitor, display.width(), display.height());

    // 创建捕获器
    let mut capturer = Capturer::new(display).expect("无法创建屏幕捕获器");
    let width = capturer.width();
    let height = capturer.height();

    // 确定输出文件名
    let output_file = args.output.unwrap_or_else(|| {
        let ext = match args.format.as_str() {
            "png-sequence" => "png",
            "mp4" => "mp4",
            _ => "gif",
        };
        format!("recording_{}.{}", Local::now().format("%Y%m%d_%H%M%S"), ext)
    });

    println!("📁 输出文件: {}", output_file);
    println!("⏳ 准备录制...\n");

    // 等待一秒让用户准备
    thread::sleep(Duration::from_secs(1));
    println!("🔴 开始录制！");

    let frame_delay = Duration::from_millis(1000 / args.fps as u64);
    let total_duration = Duration::from_secs(args.duration);
    let start_time = Instant::now();
    
    let mut frames: Vec<RgbaImage> = Vec::new();
    let mut frame_count = 0;

    // 捕获帧
    while start_time.elapsed() < total_duration {
        let frame_start = Instant::now();

        match capturer.frame() {
            Ok(buffer) => {
                // scrap 在 Windows 上返回 BGRA 格式
                // 创建 RGBA 图像
                let mut img = RgbaImage::new(width as u32, height as u32);
                
                // 正确转换 BGRA 到 RGBA
                for y in 0..height {
                    for x in 0..width {
                        let i = (y * width + x) * 4;
                        let pixel = image::Rgba([
                            buffer[i + 2], // R (从 B 位置)
                            buffer[i + 1], // G
                            buffer[i],     // B (从 R 位置)
                            255,           // A (完全不透明)
                        ]);
                        img.put_pixel(x as u32, y as u32, pixel);
                    }
                }

                frames.push(img);
                frame_count += 1;
                
                let progress = (start_time.elapsed().as_secs_f32() / total_duration.as_secs_f32() * 100.0) as u32;
                print!("\r⏺️  录制中... {}% ({} 帧)", progress, frame_count);
                std::io::Write::flush(&mut std::io::stdout()).ok();
            }
            Err(error) => {
                if error.kind() == ErrorKind::WouldBlock {
                    // 帧还未准备好，继续等待
                    thread::sleep(Duration::from_millis(1));
                    continue;
                } else {
                    eprintln!("\n❌ 捕获帧时出错: {}", error);
                    break;
                }
            }
        }

        // 控制帧率
        let elapsed = frame_start.elapsed();
        if elapsed < frame_delay {
            thread::sleep(frame_delay - elapsed);
        }
    }

    println!("\n\n✅ 录制完成！共捕获 {} 帧", frame_count);

    if frames.is_empty() {
        eprintln!("❌ 没有捕获到任何帧！");
        return;
    }

    // 根据格式保存
    match args.format.as_str() {
        "png-sequence" => {
            println!("💾 正在保存 PNG 序列...");
            if let Err(e) = save_as_png_sequence(&output_file, &frames) {
                eprintln!("❌ 保存 PNG 序列时出错: {}", e);
                return;
            }
            println!("🎉 录制成功！文件已保存至: {}", output_file);
        }
        "mp4" => {
            println!("💾 正在保存 PNG 序列并转换为 MP4...");
            if let Err(e) = save_as_mp4(&output_file, &frames, args.fps, args.keep_frames) {
                eprintln!("❌ 保存 MP4 时出错: {}", e);
                return;
            }
            println!("🎉 录制成功！MP4 文件已保存至: {}", output_file);
        }
        _ => {
            println!("💾 正在保存 GIF 文件...");
            if let Err(e) = save_as_gif(&output_file, &frames, args.fps) {
                eprintln!("❌ 保存 GIF 时出错: {}", e);
                return;
            }
            println!("🎉 录制成功！文件已保存至: {}", output_file);
        }
    }
}

fn save_as_gif(
    path: &str,
    frames: &[RgbaImage],
    fps: u32,
) -> Result<(), Box<dyn std::error::Error>> {
    if frames.is_empty() {
        return Err("没有帧可保存".into());
    }

    // 缩放以减小文件大小
    let scale = 2;
    let first_frame = &frames[0];
    let orig_width = first_frame.width();
    let orig_height = first_frame.height();
    let new_width = (orig_width / scale) as u16;
    let new_height = (orig_height / scale) as u16;

    println!("📐 缩放: {}x{} -> {}x{}", orig_width, orig_height, new_width, new_height);

    let file = BufWriter::new(File::create(path)?);
    let mut encoder = Encoder::new(file, new_width, new_height, &[])?;
    encoder.set_repeat(Repeat::Infinite)?;

    let frame_delay = (100.0 / fps as f32) as u16; // 以 1/100 秒为单位

    for (i, img) in frames.iter().enumerate() {
        // 缩放图像
        let resized = image::imageops::resize(
            img,
            new_width as u32,
            new_height as u32,
            image::imageops::FilterType::Nearest,
        );

        // 转换为 GIF 帧格式
        let mut gif_frame = Frame::from_rgba_speed(
            new_width,
            new_height,
            &mut resized.into_raw(),
            10,
        );
        gif_frame.delay = frame_delay;

        encoder.write_frame(&gif_frame)?;

        if (i + 1) % 10 == 0 || i == frames.len() - 1 {
            print!("\r💾 保存进度: {}%", (i + 1) * 100 / frames.len());
            std::io::Write::flush(&mut std::io::stdout()).ok();
        }
    }

    println!();
    Ok(())
}

fn save_as_png_sequence(
    base_path: &str,
    frames: &[RgbaImage],
) -> Result<(), Box<dyn std::error::Error>> {
    // 创建输出目录
    let dir_path = base_path.trim_end_matches(".png");
    std::fs::create_dir_all(dir_path)?;

    for (i, img) in frames.iter().enumerate() {
        let file_path = format!("{}/frame_{:05}.png", dir_path, i);
        img.save(&file_path)?;

        if (i + 1) % 10 == 0 || i == frames.len() - 1 {
            print!("\r💾 保存进度: {}%", (i + 1) * 100 / frames.len());
            std::io::Write::flush(&mut std::io::stdout()).ok();
        }
    }

    println!();
    println!("📂 PNG 序列已保存至目录: {}", dir_path);
    Ok(())
}

fn save_as_mp4(
    path: &str,
    frames: &[RgbaImage],
    fps: u32,
    keep_frames: bool,
) -> Result<(), Box<dyn std::error::Error>> {
    if frames.is_empty() {
        return Err("没有帧可保存".into());
    }

    // 创建临时目录保存 PNG 序列
    let temp_dir = format!("{}_temp_frames", path.trim_end_matches(".mp4"));
    std::fs::create_dir_all(&temp_dir)?;

    // 保存 PNG 序列
    println!("📸 正在保存临时 PNG 序列...");
    for (i, img) in frames.iter().enumerate() {
        let file_path = format!("{}/frame_{:05}.png", temp_dir, i);
        img.save(&file_path)?;

        if (i + 1) % 10 == 0 || i == frames.len() - 1 {
            print!("\r💾 保存进度: {}%", (i + 1) * 100 / frames.len());
            std::io::Write::flush(&mut std::io::stdout()).ok();
        }
    }
    println!();

    // 使用 ffmpeg 转换为 MP4
    println!("🎬 正在使用 ffmpeg 转换为 MP4...");
    let output = Command::new("ffmpeg")
        .args(&[
            "-y",                                    // 覆盖已存在的文件
            "-framerate", &fps.to_string(),         // 帧率
            "-i", &format!("{}/frame_%05d.png", temp_dir), // 输入文件模式
            "-c:v", "libx264",                      // 视频编码器
            "-preset", "medium",                     // 编码速度预设
            "-crf", "23",                           // 质量（0-51，越小越好）
            "-pix_fmt", "yuv420p",                  // 像素格式（兼容性最好）
            path,                                    // 输出文件
        ])
        .output()?;

    if !output.status.success() {
        let error_msg = String::from_utf8_lossy(&output.stderr);
        return Err(format!("ffmpeg 转换失败: {}", error_msg).into());
    }

    println!("✅ MP4 转换完成！");

    // 清理临时文件
    if !keep_frames {
        println!("🧹 正在清理临时文件...");
        std::fs::remove_dir_all(&temp_dir)?;
        println!("✅ 临时文件已清理");
    } else {
        println!("📂 临时 PNG 文件保留在: {}", temp_dir);
    }

    Ok(())
}
